import os
import numpy as np
import nibabel as nib
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import skimage.transform  # 新增导入
from scipy.ndimage import gaussian_filter, map_coordinates
import matplotlib.pyplot as plt

# 以ROI区域为截取ROI区域，并将图片缩放为cropsize*cropsize
class MRIDataset(Dataset):
    def __init__(self, root_dir, modalities=('axi', 'sag', 'cor'), max_slices=3, crop_size=128, augment=None):
        self.root_dir = root_dir
        self.modalities = modalities
        self.max_slices = max_slices
        self.crop_size = crop_size
        self.augment = augment
        self.labels = pd.read_csv(os.path.join(root_dir, 'label.csv'),
                         header=0,
                         names=['p_id', 'label']).set_index('p_id')
        self.labels['label'] = self.labels['label'].astype(int) - 1  # 转换为0-based
        assert self.labels['label'].between(0, 3).all(), "标签必须在0-3范围内"
        
        self.patient_ids = []
        for pid in self.labels.index:
            if self._check_data_complete(pid):
                self.patient_ids.append(pid)

        self.class_counts = {i: 0 for i in range(4)}  # 按0-based初始化
        for pid in self.patient_ids:
            label = self.labels.loc[pid, 'label']
            self.class_counts[label] += 1

    def _check_data_complete(self, pid):
        # 与原代码相同，保持不变
        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            if not (os.path.exists(img_path) and os.path.exists(roi_path)):
                return False
            
            try:
                img = nib.load(img_path)
                roi = nib.load(roi_path)
            except Exception as e:
                print(f"加载文件失败：{img_path} 或 {roi_path}, 错误：{e}")
                return False
            
            img_shape = img.shape[:2]
            roi_shape = roi.shape[:2]
            
            if img_shape != roi_shape:
                print(f"患者 {pid} 的 {mod} 模态图像尺寸 {img_shape} 与 ROI 尺寸 {roi_shape} 不匹配")
                return False
            
            max_allowed_ratio = 1.2
            h, w = img_shape
            if h == 0 or w == 0:
                print(f"患者 {pid} 的 {mod} 模态存在零尺寸：{img_shape}")
                return False
            current_ratio = max(h / w, w / h)
            if current_ratio > max_allowed_ratio:
                print(f"患者 {pid} 的 {mod} 模态长宽比过大：{img_shape} (比例 {current_ratio:.2f})")
                return False
            
        return True

    def _crop_around_center(self, image, center_y, center_x, is_roi=False):
        """根据给定中心坐标进行裁剪"""
        target_size = self.crop_size
        half_size = target_size // 2
        
        # 计算理论裁剪范围
        y_start = center_y - half_size
        y_end = center_y + half_size
        x_start = center_x - half_size
        x_end = center_x + half_size
        
        # 计算实际可裁剪范围
        actual_y_start = max(0, y_start)
        actual_y_end = min(image.shape[0], y_end)
        actual_x_start = max(0, x_start)
        actual_x_end = min(image.shape[1], x_end)
        
        # 执行实际裁剪
        cropped = image[actual_y_start:actual_y_end,
                        actual_x_start:actual_x_end]
        
        # 计算需要填充的边界
        pad_y_before = max(0, -y_start)
        pad_y_after = max(0, y_end - image.shape[0])
        pad_x_before = max(0, -x_start)
        pad_x_after = max(0, x_end - image.shape[1])
        
        # 使用边缘像素填充
        padded = np.pad(cropped,
                        ((pad_y_before, pad_y_after),
                         (pad_x_before, pad_x_after)),
                        mode='edge')
        
        # 对ROI进行二值化处理
        if is_roi:
            padded = (padded > 0.5).astype(np.float32)
        
        return padded.astype(np.float32)

    def __len__(self):
        return len(self.patient_ids)

    def __getitem__(self, idx):
        pid = self.patient_ids[idx]
        label = self.labels.loc[pid, 'label']
        
        multimodal_img = []
        multimodal_roi = []
        multimodal_mask = []

        for mod in self.modalities:
            img_path = os.path.join(self.root_dir, 'img', f'{pid}_T2_{mod}_000.nii.gz')
            roi_path = os.path.join(self.root_dir, 'roi', f'{pid}_T2_{mod}_roi.nii.gz')
            
            img = nib.load(img_path)
            roi = nib.load(roi_path)
            
            img_data = img.get_fdata().astype(np.float32)  # (H, W, D)
            roi_data = roi.get_fdata().astype(np.float32)
            
            # 1. 找到ROI最大的切片
            max_roi_area = -1
            max_slice_idx = -1
            for slice_idx in range(roi_data.shape[2]):
                current_area = np.sum(roi_data[..., slice_idx] > 0)
                if current_area > max_roi_area:
                    max_roi_area = current_area
                    max_slice_idx = slice_idx
            
            # 处理没有找到ROI的情况
            if max_slice_idx == -1:
                print(f"警告：患者 {pid} 的 {mod} 模态没有有效的ROI")
                img_slices = [np.zeros((self.crop_size, self.crop_size))] * 3
                roi_slices = [np.zeros((self.crop_size, self.crop_size))] * 3
                mask = [0.0] * 3
                multimodal_img.append(torch.FloatTensor(np.stack(img_slices)))
                multimodal_roi.append(torch.FloatTensor(np.stack(roi_slices)))
                multimodal_mask.append(torch.FloatTensor(mask))
                continue
                
            # 2. 计算中心坐标
            roi_max_slice = roi_data[..., max_slice_idx]
            coords = np.argwhere(roi_max_slice > 0)
            if len(coords) == 0:
                center_y = img_data.shape[0] // 2
                center_x = img_data.shape[1] // 2
            else:
                min_y, min_x = coords.min(axis=0)
                max_y, max_x = coords.max(axis=0)
                center_y = (min_y + max_y) // 2
                center_x = (min_x + max_x) // 2
            
            # 3. 生成三张切片索引（中心±1）
            selected_indices = []
            for offset in range(-1, 2):
                selected_idx = max_slice_idx + offset
                selected_indices.append(selected_idx)
            
            img_slices = []
            roi_slices = []
            mask = []
            mod_min = np.inf
            mod_max = -np.inf

            for slice_idx in range(img_data.shape[2]):
                mod_min = min(mod_min, np.min(img_data[..., slice_idx]))
                mod_max = max(mod_max, np.max(img_data[..., slice_idx]))

            for slice_idx in selected_indices:
                # 处理无效索引
                if slice_idx < 0 or slice_idx >= img_data.shape[2]:
                    img_slices.append(np.zeros((self.crop_size, self.crop_size)))
                    roi_slices.append(np.zeros((self.crop_size, self.crop_size)))
                    mask.append(0.0)
                    continue
                
                # 加载切片数据
                img_slice = img_data[..., slice_idx]
                roi_slice = roi_data[..., slice_idx]
                
                # 裁剪操作
                img_cropped = self._crop_around_center(img_slice, center_y, center_x)
                roi_cropped = self._crop_around_center(roi_slice, center_y, center_x, is_roi=True)
                
                # 归一化
                img_cropped = (img_cropped - mod_min) / (mod_max - mod_min + 1e-6)
                
                # 裁剪后进行标准化
                mean = img_cropped.mean()
                std = img_cropped.std() + 1e-6  # 防止除零
                img_cropped = (img_cropped - mean) / std
                
                # 判断是否有有效ROI
                has_roi = np.any(roi_slice > 0)
                mask.append(1.0 if has_roi else 0.0)
                
                # 数据增强
                if self.augment:
                    augmented = self.augment(image=img_cropped, mask=roi_cropped)
                    img_cropped = augmented['image']
                    roi_cropped = augmented['mask']
                
                img_slices.append(img_cropped)
                roi_slices.append(roi_cropped)
            
            multimodal_img.append(torch.FloatTensor(np.stack(img_slices)))
            multimodal_roi.append(torch.FloatTensor(np.stack(roi_slices)))
            multimodal_mask.append(torch.FloatTensor(mask))

        multimodal_img = torch.stack(multimodal_img)
        multimodal_roi = torch.stack(multimodal_roi)
        multimodal_mask = torch.stack(multimodal_mask)

        return (
            multimodal_img,   # (M, D, H, W)
            multimodal_roi,
            # multimodal_mask,
            torch.tensor(label, dtype=torch.long)
        )

    def apply_augment(self, img, roi):
        """应用同步数据增强"""
        # 弹性形变
        if np.random.rand() < 0.5:
            alpha = self.crop_size * 2
            sigma = self.crop_size * 0.08
            random_state = np.random.RandomState(None)
            
            shape = img.shape
            dx = random_state.rand(*shape) * 2 - 1
            dy = random_state.rand(*shape) * 2 - 1
            
            dx = gaussian_filter(dx, sigma) * alpha
            dy = gaussian_filter(dy, sigma) * alpha
            
            x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
            indices = (y + dy).reshape(-1, 1), (x + dx).reshape(-1, 1)
            
            img = map_coordinates(img, indices, order=1, mode='reflect').reshape(shape)
            roi = map_coordinates(roi, indices, order=0, mode='constant').reshape(shape)

        # 缩放与裁剪
        if np.random.rand() < 0.5:
            scale_factor = np.random.uniform(0.9, 1.1)
            new_size = int(self.crop_size * scale_factor)
            
            img = skimage.transform.resize(img, (new_size, new_size), order=1, preserve_range=True, anti_aliasing=True)
            roi = skimage.transform.resize(roi, (new_size, new_size), order=0, preserve_range=True, anti_aliasing=False)
            
            if new_size > self.crop_size:
                start = np.random.randint(0, new_size - self.crop_size)
                img = img[start:start+self.crop_size, start:start+self.crop_size]
                roi = roi[start:start+self.crop_size, start:start+self.crop_size]
            else:
                pad = self.crop_size - new_size
                img = np.pad(img, ((0, pad), (0, pad)), mode='constant')
                roi = np.pad(roi, ((0, pad), (0, pad)), mode='constant')

        # 旋转与平移
        if np.random.rand() < 0.5:
            angle = np.random.uniform(-10, 10)
            img = skimage.transform.rotate(img, angle, order=1, preserve_range=True, mode='reflect')
            roi = skimage.transform.rotate(roi, angle, order=0, preserve_range=True, mode='constant')

        if np.random.rand() < 0.5:
            max_offset = int(self.crop_size * 0.1)
            dy, dx = np.random.randint(-max_offset, max_offset, size=2)
            pad_size = max_offset
            img_padded = np.pad(img, pad_size, mode='reflect')
            roi_padded = np.pad(roi, pad_size, mode='constant')
            
            img = img_padded[pad_size+dy : pad_size+dy+self.crop_size, pad_size+dx : pad_size+dx+self.crop_size]
            roi = roi_padded[pad_size+dy : pad_size+dy+self.crop_size, pad_size+dx : pad_size+dx+self.crop_size]

        # 颜色扰动
        if np.random.rand() < 0.5:
            img = img * np.random.uniform(0.8, 1.2) + np.random.uniform(-0.1, 0.1)
            img = np.clip(img, img.min(), img.max())

        return img, roi

def visualize_attention(img_batch, roi_batch, labels):
    # 获取首个样本数据
    img, roi, label = img_batch[0].numpy(), roi_batch[0].numpy(), labels[0].item()
    mod_names = ['Axial', 'Sagittal', 'Coronal']
    
    # 创建图形窗口
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    plt.suptitle(f"Label: {label} | Slice: 0/{img.shape[1]-1}", fontsize=16)
    
    # 初始化切片索引
    current_slice = [0]  # 使用列表实现闭包内的可变性
    max_slice = img.shape[1] - 1
    
    # 存储图像对象
    image_plots = [[None for _ in range(3)] for _ in range(3)]
    
    # 初始化显示
    for mod in range(3):
        # 原始图像
        image_plots[0][mod] = axs[0, mod].imshow(img[mod, 0], cmap='gray')
        axs[0, mod].set_title(f"{mod_names[mod]} Image")
        
        # ROI掩膜
        image_plots[1][mod] = axs[1, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.5)
        axs[1, mod].set_title(f"{mod_names[mod]} ROI Mask")
        
        # 叠加显示
        image_plots[2][mod] = axs[2, mod].imshow(img[mod, 0], cmap='gray')
        image_plots[2][mod] = axs[2, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.3)
        axs[2, mod].set_title("Overlay")
    
    def update_display():
        """更新所有模态的显示"""
        fig.suptitle(f"Label: {label} | Slice: {current_slice[0]}/{max_slice}", fontsize=16)
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(img[mod, current_slice[0]])
            
            # 更新ROI掩膜
            image_plots[1][mod].set_data(roi[mod, current_slice[0]])
            
            # 更新叠加显示（需要重新创建以避免透明度混合问题）
            axs[2, mod].images[0].set_data(img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(roi[mod, current_slice[0]])
        
        fig.canvas.draw_idle()
    
    def on_scroll(event):
        """鼠标滚轮事件处理"""
        if event.button == 'up':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.button == 'down':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    # 绑定滚轮事件
    fig.canvas.mpl_connect('scroll_event', on_scroll)
    
    # 添加键盘事件支持
    def on_key(event):
        """键盘事件处理（左右方向键）"""
        if event.key == 'right':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.key == 'left':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    plt.tight_layout()
    plt.show()



if __name__ == "__main__":

    dataset = MRIDataset(root_dir='./datasets/pT/three_label_data/train')
    print(len(dataset))
    print(dataset.class_counts)
    print("有效患者数量:", len(dataset.patient_ids))

    loader = DataLoader(dataset, batch_size=2,shuffle=True)
    
    img_batch, roi_batch, labels = next(iter(loader))
    
    masked_img = img_batch * roi_batch

    visualize_attention(img_batch, roi_batch, labels)

    # 验证数据形状
    
    print(f"IMG   shape: {img_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
    print(f"ROI   shape: {roi_batch.shape}")
    # print(f"Batch shape: {mask_batch.shape}")  # 应该输出 (B, M=3, D=3, H=224, W=224)
